package skewreduce.lib;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.List;
import java.util.Random;
import java.util.Collection;

public class KDTree<V extends KDTreeNode<V>> {
    /*
    class Node {
        V obj;

        Node left;
        Node right;

        Node(V o) { obj = o; }

        public boolean isLeaf() { return left == null && right == null; }

        public float getCoord(int lv) {
            if ( lv == 0 ) return obj.getX();
            if ( lv == 1 ) return obj.getY();
            return obj.getZ();
        }
    }
    */
	
	public static interface Axes {
		public int getNumAxes();
		public int getField(int lv);
		public boolean contained(float[] low,float[] up,ITuple t);
	}
	
	public static final class FullAxes implements Axes {
		final int numAxes;
		public FullAxes() { this(3); }
		public FullAxes(int n) { numAxes = n; }
		@Override public int getNumAxes() { return numAxes; }
		@Override public int getField(int lv) { return lv % numAxes; }
		
		@Override public boolean contained(float[] low,float[] up,ITuple t) {
	        int j;
	        for ( j = 0; j < 3 
	                && low[j] <= t.getFloat(j) 
	                && t.getFloat(j) <= up[j]; 
	                ++j );
	        return j == 3;
		}
	}
	
	public static final class SparseAxes implements Axes {
		final BitSet skipMask = new BitSet(3);
		int numAxes;
		int[] fields;
		
		public SparseAxes() {
			numAxes = 3;
		}
		
		public SparseAxes(int... toSkip) {
			for ( int i : toSkip )
				skipMask.set(i);
			numAxes = 3 - toSkip.length;
			build();
		}
		
		public void skipAxis(int a) {
			skipMask.set(a);
			--numAxes;
		}
		
		public void build() {
			fields = new int[numAxes];
			for ( int i = 0, j = 0; i < fields.length; ++i ) {
				fields[i] = skipMask.nextClearBit(j);
				j = fields[i] + 1;
			}
		}
		
		
		@Override public int getNumAxes() {
			return numAxes;
		}
		
		@Override public int getField(int lv) {
			return fields[lv % numAxes ];
		}
		
		@Override public boolean contained(float[] low,float[] up,ITuple t) {
	        int match = 0;
	        for ( int j = 0; j < 3; ++j ) {
	        	if ( skipMask.get(j) ) continue;
	            if ( low[j] > t.getFloat(j) || t.getFloat(j) > up[j] ) break;
	            ++match;
	        }
	        return match == numAxes;
		}
	}

    final Axes axes;

    V root;
    V[] values;   // store
    Random rand;

    float[] low;
    float[] up;

    public KDTree(V[] vals) {
    	this(vals,new FullAxes(3));
    }
    
    public KDTree(V[] vals,Axes a) {
    	axes = a;
        rand = new Random();
        values = vals;
        root = buildTree(0, values.length - 1, 0);
        rand = null;

        low = new float[3];
        up = new float[3];
    }

    int partition(int left,int right,int pivot,int lv,int k) {
        float pv = values[pivot].get(lv);
        int numDupPvVals = 0;
        V tmp = values[pivot];
        values[pivot] = values[right];
        values[right] = tmp;

        int storeIdx = left;
        for ( int i = left; i < right; ++i ) {
            float v = values[i].get(lv);
            if ( v < pv ) {
                tmp = values[i];
                values[i] = values[storeIdx];
                values[storeIdx] = tmp;
                ++storeIdx;
            }
            if ( v == pv ) ++numDupPvVals;
        }
        tmp = values[storeIdx];
        values[storeIdx] = values[right];
        values[right] = tmp;
        
        // now we have all values less than pv on the left storeIdx
        // if we have any duplicate pivot values, pack them nearby storeIdx
        int lastIndex = storeIdx;
        if ( numDupPvVals > 0 ) {
            final int limit = storeIdx + numDupPvVals;

        	for ( int i = storeIdx+1; i <= right; ++i ) {
        		float v = values[i].get(lv);
        		if ( pv == v ) {
        			++lastIndex;
        			if ( i != lastIndex ) { // swap only we have to
	        			tmp = values[i];
	        			values[i] = values[lastIndex];
	        			values[lastIndex] = tmp;
        			}
        			if ( lastIndex == limit ) break; // no more duplicate pivot variables
        		}
        	}
        }
        
        // range check
        int destIndex = left + k - 1;
        if ( storeIdx <= destIndex && destIndex <= lastIndex ) {
        	return destIndex;
        }
        return ( destIndex < storeIdx ) ? storeIdx : lastIndex; // skip over duplicate part
    }

    int findKMedian(int k,int left,int right,int lv) {
        int pidx = left + rand.nextInt(right - left+1);
        int p = partition(left,right,pidx,lv,k);
        if ( p < 0 ) {
            // all got the same values! don't shuffle but pick the mid point
            return k + left - 1;
        } else {
            int rank = p - left + 1;
            if ( rank == k ) {
                return p;
            } else if ( rank < k ) {
                return findKMedian( k - rank, p+1, right, lv);
            } else {
                return findKMedian( k, left, p-1, lv);
            }
        }
    }

    private V buildTree(int from,int to,int depth) {
        if ( from > to ) {
            return null;
        } else if ( from == to ) {
            // make sure the children is null
            values[from].setLeft(null);
            values[from].setRight(null);
            return values[from];
        } else {
            int len = to - from + 1;
            //int axis = depth % 3;
            int axis = axes.getField(depth);
            int med = len >> 1;

            findKMedian(med+1,from,to,axis);

            values[from+med].setLeft( buildTree(from, from + med - 1, depth+1));
            values[from+med].setRight( buildTree(from + med + 1, to, depth+1));

            return values[from+med];
        }
    }

    public Collection<V> range(IPoint3D src,float radius) {
        List<V> result = new ArrayList<V>();
        range(src,radius,null,result);
        return result;
    }

    public int range(IPoint3D src,float radius,Collection<? super V> result) {
        return range(src,radius,null,result);
    }

    public int size() { return values.length; }

    public int range(IPoint3D src,float radius,KDTreePredicate<V> pred,Collection<? super V> result) {
        low[0] = src.getX() - radius;
        low[1] = src.getY() - radius;
        low[2] = src.getZ() - radius;
        up[0] = src.getX() + radius;
        up[1] = src.getY() + radius;
        up[2] = src.getZ() + radius;

        return rsearch(root,0,pred,result);
    }

    public int range(IPoint3D lower,IPoint3D upper,float radius,KDTreePredicate<V> pred,Collection<? super V> result) {
        low[0] = lower.getX()-radius; low[1] = lower.getY()-radius; low[2] = lower.getZ()-radius;
        up[0] = upper.getX() + radius; up[1] = upper.getY() + radius; up[2] = upper.getZ() + radius;
        return rsearch(root,0,pred,result);
    }

    private int rsearch(V t,int lev,KDTreePredicate<V> p,Collection<? super V> buf) {
        int matches = 0;

        if ( t == null || ( p != null && p.evalSkip(t) ) ) return 0;
        
        int field = axes.getField(lev);

        //if ( low[lev] <= t.get(lev) ) {
        if ( low[field] <= t.get(field) ) {
            matches += rsearch(t.getLeft(), lev + 1, p, buf);
        }

//        int j;
//        for ( j = 0; j < 3 
//                && low[j] <= t.get(j) 
//                && t.get(j) <= up[j]; 
//                ++j );
//
//        if ( j == 3 ) {
        if ( axes.contained(low, up, t) ) {
            if ( p == null || p.evalMatch(t) ) {
                ++matches;
                if ( p == null || p.evalAdd(t) ) buf.add(t);
            }
        }

        //if ( t.get(lev) <= up[lev] ) {
        if ( t.get(field) <= up[field] ) {
            matches += rsearch(t.getRight(),lev+1,p,buf);
        }

        t.leave();

        return matches;
    }

    public boolean rangeExist(IPoint3D lower,IPoint3D upper,KDTreePredicate<V> pred) {
        low[0] = lower.getX(); low[1] = lower.getY(); low[2] = lower.getZ();
        up[0] = upper.getX(); up[1] = upper.getY(); up[2] = upper.getZ();
        return rsearchExist(root,0,pred);
    }

    public boolean rangeExist(IPoint3D p,float R,KDTreePredicate<V> pred) {
        low[0] = p.getX() - R; low[1] = p.getY() - R; low[2] = p.getZ() - R;
        up[0] = p.getX() + R; up[1] = p.getY() + R; up[2] = p.getZ() + R;
        return rsearchExist(root,0,pred);
    }


    private boolean rsearchExist(V t,int lev,KDTreePredicate<V> p) {
        if ( t == null || ( p != null && p.evalSkip(t) ) )
            return false;
        
        int field = axes.getField(lev);

        if ( low[field] <= t.get(field) ) {
            if ( rsearchExist(t.getLeft(), lev+1, p) ) {
                return true;
            }
        }

//        int j;
//        for ( j = 0; j < 3 
//                && low[j] <= t.get(j) 
//                && t.get(j) <= up[j]; 
//                ++j );
//
//        if ( j == 3 ) {
        if ( axes.contained(low, up, t)) {
            if ( p == null || p.evalMatch(t) ) {
                if ( p.evalAdd(t) ) {
                    return true;
                }
            }
        }

        //if ( t.get(lev) <= up[lev] ) {
        if ( t.get(field) <= up[field] ) {
            if ( rsearchExist(t.getRight(),lev+1,p) ) {
                return true;
            }
        }

        return false;
    }
}
